#tool
import requests
from langchain_community.utilities import SerpAPIWrapper
from langchain.agents import tool
from langchain_community.vectorstores import Qdrant
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from qdrant_client import QdrantClient
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.output_parsers import JsonOutputParser,StrOutputParser
import os

os.environ["OPENAI_API_KEY"] = "sk-PzPw20RJcprceKU4BqcMtchJm3LAgqqpYIZz4zImdYV5z8eC"
os.environ["OPENAI_API_BASE"] = "http://agentapi.xyz/v1"
os.environ["SERPAPI_API_KEY"] = "6ecd33ba1d6883a806c4a678bb6e247be0528fb30deb03055aa2c7ed264af7c2"


@tool
def test():
    """Test tool"""
    return "test"


@tool
def search(query: str):
    """只有当不了解或不知道实时情况的时候才会使用搜索工具"""
    result = SerpAPIWrapper().run(query)
    return result


@tool
def get_info_from_local_db(query: str):
    """只有问到2025年的运势或者蛇年相关运势的时候，会使用这个工具"""
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"  # 选择一个预训练模型
    )
    print('embeddings:', embeddings)
    client = Qdrant(
        QdrantClient("./local_qdrant_data"),
        "local_documents",
        embeddings
    )
    print('client:', client)
    retriever = client.as_retriever()
    print(retriever)
    result = retriever.invoke(query)
    return result

@tool
def bazi_cesuan(query: str):
    """只有做八字排盘的时候会用这个工具。需要用户输入姓名和出生年月日时分，如果没有输入，则这个工具不可用"""
    url = "https://api.yuanfenju.com/index.php/v1/Bazi/paipan"

    prompt = ChatPromptTemplate.from_template(
        """
        你是一个参数查询的助手，根据用户的输入内容找出相关的参数并按json的格式返回。
        JSON的字段如下：
        -"api_key":"qKNqtEihUMfbBxgsnBUY18Pfv",
        -"name":"姓名"
        -"sex":"性别，根据名字判断，男是0，女是1"
        -"type":"历类型 0农历 1公历"
        -"year":"出生年，如1998"
        -"month":"出生月，如1"
        -"day":"出生日，如22"
        -"hours":"出生时，如12"
        -"minute":"分钟，如45"
        如果用户输入的内容找不到上面的参数，则提醒用户输入更多的内容。
        用户的输入是{input}
        """
    )

    parser = JsonOutputParser()
    prompt = prompt.partial(format_instructions=parser.get_format_instructions())


    chain = prompt | ChatOpenAI(model_name="gpt-4o-mini", temperature=0) | parser
    data = chain.invoke({"input": query})
    print('bazi_cesuan_data:', data)

    result = requests.post(url, data=data)
    if result.status_code == 200:
        print('=====八字返回数据=====')
        print(result.json())
        try:
            json = result.json()
            return json['data']['bazi_info']['bazi']
        except:
            return "八字排盘失败，请稍后再试"
    else:
        return "八字排盘失败，请稍后再试"


@tool
def yaogua():
    """只有用户需要摇卦占卜的时候才会使用这个工具"""
    url = "https://api.yuanfenju.com/index.php/v1/Zhanbu/yaogua"
    result = requests.post(url, data={"api_key": "qKNqtEihUMfbBxgsnBUY18Pfv"})
    print(result)
    if result.status_code == 200:
        print('=====摇卦返回数据=====')
        print(result.json())
        try:
            json = result.json()
            image = json['data']['image']
            print('image:', image)
            return json['data']
        except:
            return "周易占卜失败，请稍后再试"
    else:
        return "周易占卜失败，请稍后再试"


@tool
def zhougong(query: str):
    """只有用户做梦了，需要解梦的时候才会使用这个工具，用户需要输入梦到的内容，如果没有输入，提醒用户输入梦到了什么"""
    url = "https://api.yuanfenju.com/index.php/v1/Gongju/zhougong"

    prompt = ChatPromptTemplate.from_template(
        """
        你是一个总结提炼助手，以用户输入的内容提炼一个梦境关键词，梦境关键字尽量精简，最好是两个字
        比如婴儿，火车。用户的梦境内容是：{input}
        """
    )

    parser = StrOutputParser()

    chain = prompt | ChatOpenAI(model_name="gpt-4o-mini", temperature=0)
    content = chain.invoke({"input": query})
    print('zhougong_content:', content.content)
    result = requests.post(url, data={"api_key": "qKNqtEihUMfbBxgsnBUY18Pfv",
                                      "title_zhougong": content.content})
    print(result)
    if result.status_code == 200:
        print('=====周公解梦返回数据=====')
        print(result.json())
        try:
            json = result.json()
            return json['data']
        except:
            return "周公解梦失败，请稍后再试"
    else:
        return "周公解梦失败，请稍后再试"
